package com.hserver.docker.controller;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
import com.github.dockerjava.api.model.Frame;
import com.github.dockerjava.core.DefaultDockerClientConfig;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.core.DockerClientConfig;
import com.github.dockerjava.netty.NettyDockerCmdExecFactory;
import com.hserver.docker.config.ConstConfig;
import top.hserver.core.interfaces.WebSocketHandler;
import top.hserver.core.ioc.annotation.WebSocket;
import top.hserver.core.server.handlers.Ws;

import java.io.*;
import java.util.HashMap;
import java.util.Map;


/**
 * @author yiChenLove
 */
@WebSocket("/stream")
public class StreamWebSocket implements WebSocketHandler {

    private Map<String, ExecCallback> map = new HashMap<>();

    @Override
    public void onConnect(Ws ws) {
        ExecCallback connect = connect(ws);
        if (connect != null) {
            map.put(ws.getUid(), connect);
        }
    }

    @Override
    public void onMessage(Ws ws) {
        ExecCallback execCallback = map.get(ws.getUid());
        execCallback.write(ws.getMessage());

    }


    @Override
    public void disConnect(Ws ws) {
        ExecCallback execCallback = map.get(ws.getUid());
        if (execCallback != null) {
            execCallback.shutdown();
            map.remove(ws.getUid());
        }
    }


    private ExecCallback connect(Ws ws) {
        try {

            DockerClientConfig config = DefaultDockerClientConfig.createDefaultConfigBuilder()
                    .withDockerHost(ConstConfig.DOCKER_URL)
                    .build();

            NettyDockerCmdExecFactory factory = new NettyDockerCmdExecFactory();

            DockerClient dockerClient = DockerClientBuilder.getInstance(config)
                    .withDockerCmdExecFactory(factory)
                    .build();

            PipedOutputStream out = new PipedOutputStream();
            PipedInputStream stdin = new PipedInputStream(out);

            ExecCreateCmdResponse execCreateCmdResponse = dockerClient.execCreateCmd(ws.query("id"))
                    .withAttachStdout(true)
                    .withAttachStdin(true)
                    .withTty(true)
                    .withCmd("/bin/sh").exec();

            ExecCallback execCallback = new ExecCallback(ws, out, dockerClient);
            dockerClient.execStartCmd(execCreateCmdResponse.getId())
                    .withStdIn(stdin)
                    .withTty(false)
                    .exec(execCallback);
            return execCallback;

        } catch (Exception e) {
            e.printStackTrace();

        }
        return null;
    }

    public static class ExecCallback extends ResultCallback.Adapter<Frame> {

        private Ws ws;

        private PipedOutputStream out;

        private DockerClient dockerClient;

        ExecCallback(Ws ws, PipedOutputStream out, DockerClient dockerClient) {
            this.ws = ws;
            this.out = out;
            this.dockerClient = dockerClient;
        }

        @Override
        public void onNext(Frame item) {
            String s = new String(item.getPayload());
            //包含\n 替换为 \r\n
            if (!s.contains("\r\n") && s.contains("\n")) {
                s = s.replaceAll("\n", "\r\n");
            }
            ws.send(s);
        }

        public void write(String cmd) {
            try {
                out.write(cmd.getBytes("utf-8"));
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        public void shutdown() {
            try {
                this.close();
            } catch (IOException e) {
                e.printStackTrace();
            }

            try {
                dockerClient.close();
            } catch (IOException e) {
                e.printStackTrace();
            }

        }

    }
}
